# --- Load required packages ---
library(tidyverse)
library(rdrobust)

# --- Load dataset (daily total & grid posts) ---
load("data/daily_first_outbreak.RData")
load("data/daily_first_outbreak_robust.RData")

# --- Define cutoff date (lockdown date for RDD) ---
given_date = as.Date("2020-01-23")

# --- Compute day difference relative to cutoff ---
daily_first_outbreak$difference <- as.integer(daily_first_outbreak$date - given_date)

# --- Initialize 'week' variable ---
daily_first_outbreak$week <- NA

# --- Create a weekly index around the cutoff date ---
#   week = 0  : from -3 to +3 days around Jan 23 (the treatment week)
#   week < 0  : weeks before lockdown (negative)
#   week > 0  : weeks after lockdown (positive)
daily_first_outbreak <- daily_first_outbreak |>
  mutate(
    week = case_when(
      difference >= -3 & difference <= 3 ~ 0,
      difference < -3 ~ floor((difference + 3) / 7),
      difference > 3 ~ ceiling((difference - 3) / 7),
      TRUE ~ NA_integer_
    ),
    # Replace missing daily counts with zeros
    total_count = replace_na(total_count, 0),
    grid_count  = replace_na(grid_count, 0)
  )

# --- Aggregate from daily to weekly level ---
#   Sum total and grid article counts by relative week
weekly_count_first_outbreak <- daily_first_outbreak |>
  group_by(week) |>
  summarise(
    weekly_total = sum(total_count, na.rm = TRUE),
    weekly_grid  = sum(grid_count,  na.rm = TRUE)
  )

# --- Compute 'other' articles (non-grid) ---
weekly_count_first_outbreak$others <- weekly_count_first_outbreak$weekly_total -
  weekly_count_first_outbreak$weekly_grid

# control the number of other articles (column 1)
main_control <- rdrobust(weekly_count_first_outbreak$weekly_grid,
                         weekly_count_first_outbreak$week,
                         covs = weekly_count_first_outbreak$others,
                         p = 1,
                         kernel = "triangular",
                         bwselect = "mserd",
                         all   = TRUE )
summary(main_control)

# robustness check (excluding the Jiangsu province)
daily_first_outbreak2$difference <- as.integer(daily_first_outbreak2$date - given_date)
daily_first_outbreak2$week <- NA

daily_first_outbreak2 <- daily_first_outbreak2 |>
  mutate(week = case_when(
    difference >= -3 & difference <= 3 ~ 0,
    difference < -3 ~ floor((difference + 3) / 7),
    difference > 3 ~ ceiling((difference - 3) / 7),
    TRUE ~ NA_integer_
  ))|>
  mutate(week = replace_na(week, 0),
         total_count = replace_na(total_count, 0),
         grid_count = replace_na(grid_count, 0))

weekly_count_first_outbreak2 <- daily_first_outbreak2 |>
  group_by(week) |>
  summarise(weekly_total = sum(total_count, na.rm = TRUE),
            weekly_grid = sum(grid_count, na.rm = TRUE))

weekly_count_first_outbreak2$others <- weekly_count_first_outbreak2$weekly_total - weekly_count_first_outbreak2$weekly_grid

# RDD robustness check
main2_control <- rdrobust(weekly_count_first_outbreak2$weekly_grid,
                          weekly_count_first_outbreak2$week,
                          p = 1,
                          kernel = "triangular",
                          covs = weekly_count_first_outbreak2$others,
                          bwselect = "mserd",
                          all   = TRUE )

summary(main2_control)

# plot
bws_vs <- main_control$bws[1]
week_before <- subset(weekly_count_first_outbreak, week <= 0 & week >= -bws_vs)
week_after <- subset(weekly_count_first_outbreak, week >= 0 & week <= bws_vs )

bws_vs2 <- main2_control$bws[1]
week_before.2 <- subset(weekly_count_first_outbreak2, week <= 0 & week >= -bws_vs2)
week_after.2 <- subset(weekly_count_first_outbreak2, week >= 0 & week <= bws_vs2)

figure1 <- daily_first_outbreak |>
  ggplot(aes(x = date, y = grid_count)) +
  geom_point(aes(shape = factor(ifelse(date < as.Date("2020-01-23"), "Hollow", "Solid"))), size = 2, color = "black") +
  labs(x = "", y = "Daily Grid-related News Articles", title = "") +   
  geom_vline(xintercept = as.Date("2020-01-23"), linetype = "dashed", color = "#0033FF", linewidth = 1.5) +
  theme_bw() +
  theme(plot.title = element_text(size = 17.8, lineheight = 3, face = "bold"),
        axis.text.x = element_text(size = 16),
        axis.title.x = element_text(size = 16),
        axis.title.y = element_text(size = 16),
        axis.text.y = element_text(size = 16),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(), 
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)
  ) +
  # Text label for "Wuhan Lockdown"
  annotate("text", x = as.Date("2019-10-15"), y = 36, label = "Wuhan Lockdown", size = 7, color = "black", hjust = 1) +
  annotate("segment", x = as.Date("2019-10-20"), xend = as.Date("2020-01-18"), y = 36, yend = 36, 
           colour = "black", linewidth = 0.8, arrow = arrow(length = unit(0.2, "cm"))) +
  scale_x_date(date_breaks = "1 year", date_labels = "%Y", limits = c(as.Date("2018-12-31"), as.Date("2020-12-31"))) +
  scale_shape_manual(values = c("Hollow" = 1, "Solid" = 16)) +  
  guides(shape = "none")  

figure1

figureA5 <- ggplot() +
  geom_smooth(data = week_before, aes(x = week, y = weekly_grid), method = "lm", color = "black", se = FALSE) +
  geom_point(data = week_before, aes(x = week, y = weekly_grid), color = "black", size = 5, shape = 1) +  
  geom_smooth(data = week_after, aes(x = week, y = weekly_grid), method = "lm", color = "black", se = FALSE) +
  geom_point(data = week_after, aes(x = week, y = weekly_grid), color = "black", size = 5, shape = 19) +  
  labs(x = "", y = "Number of Grid-related Articles", title = "") +  
  geom_vline(xintercept = 0, linetype = "dashed", color = "#0033FF", linewidth = 1.5) +
  annotate("text", x = -4, y = 185, label = "Wuhan Lockdown", size = 6, color = "black", hjust = 1) +
  annotate("segment", x = -3.8, xend = -0.2, y = 185, yend = 185, 
           color = "black", linewidth = 0.8, arrow = arrow(length = unit(0.2, "cm"))) +
  theme_bw() +
  theme(plot.title = element_text(size = 18, lineheight = 3, face = "bold"),
        axis.text.x = element_text(size = 16),
        axis.title.x = element_text(size = 16),
        axis.title.y = element_text(size = 16),
        axis.text.y = element_text(size = 16),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(), 
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)
  ) +
  scale_x_continuous(breaks = c(-10, -5, 0, 5, 10)) +
  scale_y_continuous(breaks = c(0, 50, 100, 150, 200), limits = c(-10, 220))


figureA5

ggsave(filename = "plots/FIGURE1.png", plot = figure1, width = 12.5, height = 5.4, units = "in", dpi = 800)
ggsave(filename = "plots/FIGUREA5.png", plot = figureA5, width = 12.4, height = 5.4, units = "in", dpi = 800)
